import numpy as np
tp1 = "mtt_result_layers_tp1/"
tp2 = "mtt_result_layers_tp2/"
if 0:
  a  = np.load(f'{tp1}/layer.0.0.q-after-rotrary.npy')
  a0 = np.load(f'{tp2}/layer.0.0.q-after-rotrary.npy')
  a1 = np.load(f'{tp2}/layer.0.1.q-after-rotrary.npy')
  a2 = np.concatenate((a0, a1), axis=1)
  print(a.shape, a0.shape, a1.shape)
  print("a")
  print(a)
  print("a0")
  print(a0)
  print("a1")
  print(a1)
  print(a - a2)

if 0:
  a  = np.load(f'{tp1}/layer.0.0.scale_dot_attn_out.npy')
  a0 = np.load(f'{tp2}/layer.0.0.scale_dot_attn_out.npy')
  a1 = np.load(f'{tp2}/layer.0.1.scale_dot_attn_out.npy')
  a2 = np.concatenate((a0, a1), axis=1)
  # np.set_printoptions(threshold=np.inf, linewidth = 24, edgeitems = 5)
  # np.set_printoptions(linewidth = 100, edgeitems = 5)
  print(a.shape, a0.shape, a1.shape)
  # print("a")
  # print(a)
  # print("a0")
  # print(a0)
  # print("a1")
  # print(a1)
  # print("a2")
  # print(a2)
  print(a - a2)
  # diff = np.where(a != a2)
  # np.set_printoptions(threshold=np.inf)
  # print(diff)
  a  = np.load(f'{tp1}/layer.0.0.task.npy')
  a0 = np.load(f'{tp2}/layer.0.0.task.npy')
  a1 = np.load(f'{tp2}/layer.0.1.task.npy')
  print(a.shape, a0.shape, a1.shape)
  print("a")
  print(a)
  print("a0")
  print(a0)
  print("a1")
  print(a1)

if 1:
  a  = np.load(f'{tp1}/layer.0.0.attn_blk_add_residual.npy')
  a0 = np.load(f'{tp2}/layer.0.0.attn_blk_add_residual.npy')
  a1 = np.load(f'{tp2}/layer.0.1.attn_blk_add_residual.npy')
  print(a.shape, a0.shape, a1.shape)
  print("a")
  print(a)
  print("a0")
  print(a0)
  print("a1")
  print(a1)
  print("diff")
  print(a - (a1 + a0))
  print(np.max(a - (a1 + a0)))

if 0:
  a  = np.load(f'{tp1}/layer.0.0.ffn_out.npy')
  a0 = np.load(f'{tp2}/layer.0.0.ffn_out.npy')
  a1 = np.load(f'{tp2}/layer.0.1.ffn_out.npy')
  print(a.shape, a0.shape, a1.shape)
  print("a")
  print(a)
  print("a0")
  print(a0)
  print("a1")
  print(a1)

# kv_cache shape: num_page, _2, layers, local_kv_head_num, tokens_per_page, head_dim
if 0:
  a  = np.load(f'{tp1}/31.0.kv_cache.npy')
  a0 = np.load(f'{tp2}/31.0.kv_cache.npy')
  a1 = np.load(f'{tp2}/31.1.kv_cache.npy')
  num_page, _2, layers, local_kv_head_num, tokens_per_page, head_dim = a.shape
  slot_mapping1 = [22976, 22977, 22978, 22979, 22980, 22981, 22982, 22983, 22984, 22985]
  slot_mapping2 = [70848, 70849, 70850, 70851, 70852, 70853, 70854, 70855, 70856, 70857]
  tokens = len(slot_mapping1)
  # kv1 = np.zeros([tokens, _2, layers, local_kv_head_num, head_dim], dtype = a.dtype)
  # for i in range(tokens):
  #   s = slot_mapping1[i]
  #   kv1[i] = a[s // tokens_per_page, :, :, :, s % tokens_per_page, :]

  kv  =  a[[s // tokens_per_page for s in slot_mapping1], :, :, :, [s % tokens_per_page for s in slot_mapping1], :]
  kv0 = a0[[s // tokens_per_page for s in slot_mapping2], :, :, :, [s % tokens_per_page for s in slot_mapping2], :]
  kv1 = a1[[s // tokens_per_page for s in slot_mapping2], :, :, :, [s % tokens_per_page for s in slot_mapping2], :]
  kv2 = np.concatenate((kv0, kv1), axis=3)
  print("tokens, _2, layers, kv_head_num, head_dim", kv1.shape, kv2.shape)
  print("layers, _2, tokens, kv_head_num, head_dim 2, 1, 0, 3, 4")
  print((kv- kv2).transpose(2, 1, 0, 3, 4))
  diff = np.where(kv != kv2)
  print(len(diff[0]))
